from operator import itemgetter
import math
import Numbers
import itertools
from numpy import matrix,array

class DisjointSet(dict):
    def add(self, item):
        self[item] = item

    def find(self, item):
        parent = self[item]

        while self[parent] != parent:
            parent = self[parent]

        self[item] = parent
        return parent

    def union(self, item1, item2):
        self[item2] = self[item1]

def kruskal( nodes, edges ):
    '''
    nodes = list( "ABCDEFG" )
    edges = [ ("A", "B", 7), ("A", "D", 5),
    ("B", "C", 8), ("B", "D", 9), ("B", "E", 7),
    ("C", "E", 5),
    ("D", "E", 15), ("D", "F", 6),
    ("E", "F", 8), ("E", "G", 9),
    ("F", "G", 11)]
    '''
    forest = DisjointSet()
    mst = []
    for n in nodes:
        forest.add( n )

    sz = len(nodes) - 1

    for e in sorted( edges, key=itemgetter( 2 ) ):
        n1, n2, _ = e
        t1 = forest.find(n1)
        t2 = forest.find(n2)
        if t1 != t2:
            mst.append(e)
            sz -= 1
            if sz == 0:
                return mst

            forest.union(t1, t2)

def floyd(f, x0):
    # The main phase of the algorithm, finding a repetition x_mu = x_2mu
    # The hare moves twice as quickly as the tortoise
    # Eventually they will both be inside the cycle
    # and the distance between them will increase by 1 until
    # it is divisible by the length of the cycle.
    tortoise = f(x0) # f(x0) is the element/node next to x0.
    hare = f(f(x0))
    while tortoise != hare:
        tortoise = f(tortoise)
        hare = f(f(hare))

    # at this point the position of tortoise which is the distance between
    # hare and tortoise is divisible by the length of the cycle.
    # so hare moving in circle and tortoise (set to x0) moving towards
    # the circle will intersect at the beginning of the circle.

    # Find the position of the first repetition of length mu
    # The hare and tortoise move at the same speeds
    mu = 0
    tortoise = x0
    while tortoise != hare:
        tortoise = f(tortoise)
        hare = f(hare)
        mu += 1

    # Find the length of the shortest cycle starting from x_mu
    # The hare moves while the tortoise stays still
    lam = 1
    hare = f(tortoise)
    while tortoise != hare:
        hare = f(hare)
        lam += 1

    return lam, mu

def quadratic_equation(a,b,c):
    d = b*b - 4*a*c
    if d < 0:
        return None
    elif d == 0:
         x = (-b + math.sqrt(b*b - 4*a*c)) / (2*a)
         return (x,x)
    else:
        x1 = (-b + math.sqrt(b*b - 4*a*c)) / (2*a)
        x2 = (-b - math.sqrt(b**2-4*a*c))  / (2*a)
        return (x1,x2)



def primitive_pyth_triples(max=None):
  '''generate all primative triples such that
  the sum is less than or equal to max'''
  u=matrix([[1,2,2],[-2,-1,-2],[2,2,3]])
  a=matrix([[1,2,2],[2,1,2],[2,2,3]])
  d=matrix([[-1,-2,-2],[2,1,2],[2,2,3]])
  m=[array([3,4,5])]
  while m:
    for i in m:
      yield i
    g=((i*j).getA1() for i in m for j in (u,a,d))
    m=[i for i in g if max is None or sum(i)<=max]

def all_pyth_triples(max):
  '''generate all triples such that
  the sum is less than or equal to max'''
  for i in primitive_pyth_triples(max):
    ret=i[:]
    while sum(ret)<=max:
      yield ret
      ret=ret+i

